'''
 * @Author: Benjay·Shaw
 * @Date: 2024-10-30 17:48:20
 * @LastEditors: Benjay·Shaw
 * @LastEditTime: 2024-10-31 14:56:39
 * @Description: UNetR50 and UNetR101
'''
import paddle
from utils.common_module import *
from functools import partial
relu_func = partial(paddle.nn.functional.relu, inplace=True)


class UNetR50(paddle.nn.Layer):

    def __init__(self, num_classes=1):
        super(UNetR50, self).__init__()
        filters = [256, 512, 1024, 2048]
        resnet = paddle.vision.models.resnet50(pretrained=True)
        self.firstconv = resnet.conv1
        self.firstbn = resnet.bn1
        self.firstrelu = resnet.relu
        self.firstmaxpool = resnet.maxpool
        self.encoder1 = resnet.layer1
        self.encoder2 = resnet.layer2
        self.encoder3 = resnet.layer3
        self.encoder4 = resnet.layer4
        self.dblock = DBlockMoreDilate(filters[3])
        self.spp = SPPBlock(filters[3])
        self.decoder4 = DecoderBlock(filters[3] + 3, filters[2])
        self.decoder3 = DecoderBlock(filters[2], filters[1])
        self.decoder2 = DecoderBlock(filters[1], filters[0])
        self.decoder1 = DecoderBlock(filters[0], filters[0])
        self.finaldeconv1 = paddle.nn.Conv2DTranspose(in_channels=filters[0
            ], out_channels=32, kernel_size=4, stride=2, padding=1)
        self.finalrelu1 = relu_func
        self.finalconv2 = paddle.nn.Conv2D(in_channels=32, out_channels=32,
            kernel_size=3, padding=1)
        self.finalrelu2 = relu_func
        self.finalconv3 = paddle.nn.Conv2D(in_channels=32, out_channels=
            num_classes, kernel_size=3, padding=1)

    def forward(self, x):
        x = self.firstconv(x)
        x = self.firstbn(x)
        x = self.firstrelu(x)
        x = self.firstmaxpool(x)
        e1 = self.encoder1(x)
        e2 = self.encoder2(e1)
        e3 = self.encoder3(e2)
        e4 = self.encoder4(e3)
        e4 = self.dblock(e4)
        d4 = self.decoder4(e4)
        d3 = self.decoder3(d4)
        d2 = self.decoder2(d3)
        d1 = self.decoder1(d2)
        out = self.finaldeconv1(d1)
        out = self.finalrelu1(out)
        out = self.finalconv2(out)
        out = self.finalrelu2(out)
        out = self.finalconv3(out)
        return paddle.nn.functional.sigmoid(x=out)


class UNetR101(paddle.nn.Layer):

    def __init__(self, num_classes=1):
        super(UNetR101, self).__init__()
        filters = [256, 512, 1024, 2048]
        resnet = paddle.vision.models.resnet101(pretrained=True)
        self.firstconv = resnet.conv1
        self.firstbn = resnet.bn1
        self.firstrelu = resnet.relu
        self.firstmaxpool = resnet.maxpool
        self.encoder1 = resnet.layer1
        self.encoder2 = resnet.layer2
        self.encoder3 = resnet.layer3
        self.encoder4 = resnet.layer4
        self.dblock = DBlockMoreDilate(2048)
        self.decoder4 = DecoderBlock(filters[3], filters[2])
        self.decoder3 = DecoderBlock(filters[2], filters[1])
        self.decoder2 = DecoderBlock(filters[1], filters[0])
        self.decoder1 = DecoderBlock(filters[0], filters[0])
        self.finaldeconv1 = paddle.nn.Conv2DTranspose(in_channels=filters[0
            ], out_channels=32, kernel_size=4, stride=2, padding=1)
        self.finalrelu1 = relu_func
        self.finalconv2 = paddle.nn.Conv2D(in_channels=32, out_channels=32,
            kernel_size=3, padding=1)
        self.finalrelu2 = relu_func
        self.finalconv3 = paddle.nn.Conv2D(in_channels=32, out_channels=
            num_classes, kernel_size=3, padding=1)

    def forward(self, x):
        x = self.firstconv(x)
        x = self.firstbn(x)
        x = self.firstrelu(x)
        x = self.firstmaxpool(x)
        e1 = self.encoder1(x)
        e2 = self.encoder2(e1)
        e3 = self.encoder3(e2)
        e4 = self.encoder4(e3)
        e4 = self.dblock(e4)
        d4 = self.decoder4(e4) + e3
        d3 = self.decoder3(d4) + e2
        d2 = self.decoder2(d3) + e1
        d1 = self.decoder1(d2)
        out = self.finaldeconv1(d1)
        out = self.finalrelu1(out)
        out = self.finalconv2(out)
        out = self.finalrelu2(out)
        out = self.finalconv3(out)
        return paddle.nn.functional.sigmoid(x=out)
